# -*- coding: utf-8 -*-
# @Time    : 2021/4/19
# @File    : utils.py
import av
import numpy as np
import torch
from torchvision import transforms

# 用于预先训练的PyTorch模型的平均值和标准偏差
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])


# 提取帧，用于视频渲染
def extract_frames(video_path):
	""" Extracts frames from video """
	frames = []
	video = av.open(video_path)
	for frame in video.decode(0):
		yield frame.to_image()


def gram_matrix(y):
	""" 返回y的 gram 矩阵(用于计算风格损失) """
	(b, c, h, w) = y.size()
	features = y.view(b, c, w * h)
	features_t = features.transpose(1, 2)
	gram = features.bmm(features_t) / (c * h * w)
	return gram


# size (sequence or int): Desired output size. If size is a sequence like
# (h, w), output size will be matched to this. If size is an int,
# smaller edge of the image will be matched to this number.
# i.e, if height > width, then image will be rescaled to
# (size * height / width, size)

def train_transform(image_size):
	""" 图像的变换 """
	transform = transforms.Compose(
		[
			transforms.Resize(image_size),
			transforms.RandomCrop(image_size),
			transforms.ToTensor(),
			transforms.Normalize(mean, std),
		]
	)
	return transform


def style_transform(image_size):
	""" 风格图像的变换 """
	resize = [transforms.Resize(image_size)] if image_size else []
	transform = transforms.Compose(resize + [transforms.ToTensor(), transforms.Normalize(mean, std)])
	return transform


def denormalize(tensors):
	""" 使用均值和标准差去规格化图像张量 """
	for channel in range(3):
		tensors[:, channel].mul_(std[channel]).add_(mean[channel])
	return tensors


def deprocess(image_tensor):
	""" 对图像张量进行去归一化和重新缩放 """
	image_tensor = denormalize(image_tensor)[0]
	image_tensor *= 255
	image_np = torch.clamp(image_tensor, 0, 255).cpu().numpy().astype(np.uint8)
	image_np = image_np.transpose(1, 2, 0)
	return image_np
